{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "47120857-b2aa-4a36-8733-e04776ca7a80",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os,sys,trimesh\n",
    "import numpy as np\n",
    "sys.path.append('/home/anpei/Code/nglod/sdf-net')\n",
    "\n",
    "os.environ['PYOPENGL_PLATFORM'] = 'egl'\n",
    "os.environ['MESA_GL_VERSION_OVERRIDE'] = '3.3'\n",
    "os.environ['MESA_GLSL_VERSION_OVERRIDE'] = '330'\n",
    "\n",
    "import torch\n",
    "from lib.torchgp import load_obj, point_sample, sample_surface, compute_sdf, normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0c8bf6fa-7295-40cb-aaaa-3cfbf95350a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_obj(path):\n",
    "    mesh = trimesh.load(path)\n",
    "    mesh.vertices = normalize(mesh.vertices)\n",
    "    return mesh\n",
    "    \n",
    "def normalize(V):\n",
    "\n",
    "    # Normalize mesh\n",
    "    V_max = np.max(V, axis=0)\n",
    "    V_min = np.min(V, axis=0)\n",
    "    V_center = (V_max + V_min) / 2.\n",
    "    V = V - V_center\n",
    "\n",
    "    # Find the max distance to origin\n",
    "    max_dist = np.sqrt(np.max(np.sum(V**2, axis=-1)))\n",
    "    V_scale = 1. / max_dist\n",
    "    V *= V_scale\n",
    "    return V\n",
    "\n",
    "\n",
    "def resample(V,F,num_samples, chunk=10000, sample_mode=['rand', 'near', 'near', 'near', 'near']):\n",
    "    \"\"\"Resample SDF samples.\"\"\"\n",
    "\n",
    "    points, sdfs = [],[]\n",
    "    for _ in range(num_samples//chunk):\n",
    "\n",
    "        pts = point_sample(V, F, sample_mode, chunk)\n",
    "        sdf = compute_sdf(V, F, pts.cuda())   \n",
    "        points.append(pts.cpu())\n",
    "        sdfs.append(sdf.cpu())\n",
    "    return torch.cat(points), torch.cat(sdfs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b74c1144-f228-473e-abe3-059b9ae19d9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for item in ['statuette.ply','dragon.ply','armadillo.obj','lucy.ply']:#'statuette.ply','dragon.ply','armadillo.obj'\n",
    "    dataset_path = f'/home/anpei/Dataset/mesh/obj/{item}'\n",
    "    mesh = load_obj(dataset_path)\n",
    "    \n",
    "    f = SDF(mesh.vertices, mesh.faces);\n",
    "\n",
    "    V = torch.from_numpy(mesh.vertices).float().cuda()\n",
    "    F = torch.from_numpy(mesh.faces).cuda()\n",
    "    \n",
    "    pts_train, _ = resample(V, F, num_samples=int(8*1024*1024/5))\n",
    "    sdf_train = f(pts_train.numpy())\n",
    "    \n",
    "    pts_test, _ = resample(V, F, num_samples=int(16*1024*1024//5))\n",
    "    sdf_test = f(pts_test.numpy())\n",
    "\n",
    "    np.save(f'/home/anpei/Dataset/mesh/obj/sdf/{item[:-4]}_8M',{'points_train':pts_train.numpy(),'sdfs_train':sdf_train, \\\n",
    "                                                            'points_test':pts_test.numpy(),'sdfs_test':sdf_test})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dafbf526-e79a-47a9-bf7b-6da70497c3af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
